import os
import librosa
import samplerate
import numpy as np
import soundfile as sf
from scipy import signal
import matplotlib.pyplot as plt

def get_all_file(folder_path):
    file_paths = []

    for root, dirs, files in os.walk(folder_path):
        # 遍历当前文件夹下的所有文件
        for file_name in files:
            file_path = os.path.join(root, file_name)
            file_paths.append(file_path)
    return file_paths

def get_all_file_name(file_paths):
    file_names = []
    for root, dirs, files in os.walk(folder_path):
        for file_name in files:
            file_name = file_name.strip('.wav')
            file_names.append(file_name)
    return file_names

def plot_signal(audio_data, title=None):
    plt.figure(figsize=(12, 3.5), dpi=300)
    plt.plot(audio_data, linewidth=1)
    plt.title(title,fontsize = 16)
    plt.tick_params(labelsize=12)
    plt.grid()
    plt.show()

def band_pass_filter(original_signal, order, fc1,fc2, fs):
    '''
    中值滤波器
    :param original_signal: 音频数据
    :param order: 滤波器阶数
    :param fc1: 截止频率
    :param fc2: 截止频率
    :param fs: 音频采样率
    :return: 滤波后的音频数据
    '''
    b, a = signal.butter(N=order, Wn=[2*fc1/fs,2*fc2/fs], btype='bandpass')
    new_signal = signal.lfilter(b, a, original_signal)
    return new_signal

folder_path = r".\four_audios"
file_names = get_all_file_name(folder_path)
file_path = get_all_file(folder_path)
print(len(file_path),len(file_names))

for file,file_name in zip(file_path,file_names):

    audio_data, sr = librosa.load(file, sr=2000)
    # plot_signal(audio_data, title='Initial Audio')

    audio_data = band_pass_filter(audio_data, 2, 25, 400, sr)
    # plot_signal(audio_data, title='After Filter')

    down_sample_audio_data = samplerate.resample(audio_data.T, 1000 / sr, converter_type='sinc_best').T
    # plot_signal(down_sample_audio_data, title='Down_sampled')

    down_sample_audio_data = down_sample_audio_data / np.max(np.abs(down_sample_audio_data))
    # plot_signal(down_sample_audio_data, title='Normalized')
    print(2)
    time_segment = 2
    overlap = 1000
    change_sr = 1000
    cut_point = int(change_sr * (time_segment - overlap / 1000))
    cut_num = len(down_sample_audio_data)//cut_point - 1
    for i in range(cut_num):
        start = cut_point*i
        end = cut_point*i+time_segment*change_sr
        segment = audio_data[start:end]
        sf.write(f'./four_audios_cut/{file_name}_{i}.wav', segment, change_sr)


